"""
SQL语法分析器 (Parser)
使用递归下降解析法实现SQL语句的语法分析，生成抽象语法树(AST)
"""

from typing import List, Optional, Union, Any
from enum import Enum
from lexer import Token, TokenType, Lexer


class ASTNode:
    """抽象语法树节点基类"""
    pass


class Statement(ASTNode):
    """语句节点"""
    pass


class Expression(ASTNode):
    """表达式节点"""
    pass


class SelectStatement(Statement):
    """SELECT语句"""
    def __init__(self, select_list: List[Expression], from_clause: 'FromClause',
                 where_clause: Optional[Expression] = None):
        self.select_list = select_list
        self.from_clause = from_clause
        self.where_clause = where_clause


class InsertStatement(Statement):
    """INSERT语句"""
    def __init__(self, table_name: str, columns: Optional[List[str]] = None,
                 values: List[List[Expression]] = None):
        self.table_name = table_name
        self.columns = columns
        self.values = values


class UpdateStatement(Statement):
    """UPDATE语句"""
    def __init__(self, table_name: str, assignments: List['Assignment'],
                 where_clause: Optional[Expression] = None):
        self.table_name = table_name
        self.assignments = assignments
        self.where_clause = where_clause


class DeleteStatement(Statement):
    """DELETE语句"""
    def __init__(self, table_name: str, where_clause: Optional[Expression] = None):
        self.table_name = table_name
        self.where_clause = where_clause


class CreateTableStatement(Statement):
    """CREATE TABLE语句"""
    def __init__(self, table_name: str, columns: List['ColumnDefinition']):
        self.table_name = table_name
        self.columns = columns


class DropTableStatement(Statement):
    """DROP TABLE语句"""
    def __init__(self, table_name: str):
        self.table_name = table_name


class FromClause(ASTNode):
    """FROM子句"""
    def __init__(self, table_name: str, alias: Optional[str] = None):
        self.table_name = table_name
        self.alias = alias


class ColumnDefinition(ASTNode):
    """列定义"""
    def __init__(self, column_name: str, data_type: str, constraints: List[str] = None):
        self.column_name = column_name
        self.data_type = data_type
        self.constraints = constraints or []


class Assignment(ASTNode):
    """赋值操作"""
    def __init__(self, column_name: str, expression: Expression):
        self.column_name = column_name
        self.expression = expression


class BinaryExpression(Expression):
    """二元表达式"""
    def __init__(self, left: Expression, operator: str, right: Expression):
        self.left = left
        self.operator = operator
        self.right = right


class UnaryExpression(Expression):
    """一元表达式"""
    def __init__(self, operator: str, operand: Expression):
        self.operator = operator
        self.operand = operand


class ColumnReference(Expression):
    """列引用"""
    def __init__(self, column_name: str, table_name: Optional[str] = None):
        self.column_name = column_name
        self.table_name = table_name


class Literal(Expression):
    """字面量"""
    def __init__(self, value: Union[str, int, float], value_type: str):
        self.value = value
        self.value_type = value_type


class FunctionCall(Expression):
    """函数调用"""
    def __init__(self, function_name: str, arguments: List[Expression]):
        self.function_name = function_name
        self.arguments = arguments


class Parser:
    """SQL语法分析器"""

    def __init__(self, tokens: List[Token]):
        self.tokens = tokens
        self.position = 0
        self.current_token = self.tokens[0] if tokens else None

    def advance(self):
        """前移到下一个词法单元"""
        self.position += 1
        if self.position >= len(self.tokens):
            self.current_token = None
        else:
            self.current_token = self.tokens[self.position]

    def peek(self) -> Optional[Token]:
        """查看下一个词法单元"""
        peek_pos = self.position + 1
        if peek_pos >= len(self.tokens):
            return None
        return self.tokens[peek_pos]

    def expect(self, token_type: TokenType) -> Token:
        """期望特定的词法单元类型"""
        if self.current_token and self.current_token.type == token_type:
            token = self.current_token
            self.advance()
            return token
        else:
            expected = token_type.value
            actual = self.current_token.type.value if self.current_token else "EOF"
            raise SyntaxError(f"Expected {expected}, but got {actual} at line {self.current_token.line if self.current_token else 'EOF'}")

    def parse(self) -> List[Statement]:
        """解析SQL语句"""
        statements = []

        while self.current_token and self.current_token.type != TokenType.EOF:
            statement = self.parse_statement()
            statements.append(statement)

            # 跳过分号
            if self.current_token and self.current_token.type == TokenType.SEMICOLON:
                self.advance()

        return statements

    def parse_statement(self) -> Statement:
        """解析语句"""
        if not self.current_token:
            raise SyntaxError("Unexpected end of input")

        if self.current_token.type == TokenType.SELECT:
            return self.parse_select_statement()
        elif self.current_token.type == TokenType.INSERT:
            return self.parse_insert_statement()
        elif self.current_token.type == TokenType.UPDATE:
            return self.parse_update_statement()
        elif self.current_token.type == TokenType.DELETE:
            return self.parse_delete_statement()
        elif self.current_token.type == TokenType.CREATE:
            return self.parse_create_statement()
        elif self.current_token.type == TokenType.DROP:
            return self.parse_drop_statement()
        else:
            raise SyntaxError(f"Unexpected token: {self.current_token.type.value}")

    def parse_select_statement(self) -> SelectStatement:
        """解析SELECT语句"""
        self.expect(TokenType.SELECT)

        # 解析选择列表
        select_list = self.parse_select_list()

        # 解析FROM子句
        self.expect(TokenType.FROM)
        from_clause = self.parse_from_clause()

        # 可选的WHERE子句
        where_clause = None
        if self.current_token and self.current_token.type == TokenType.WHERE:
            self.advance()
            where_clause = self.parse_expression()

        return SelectStatement(select_list, from_clause, where_clause)

    def parse_select_list(self) -> List[Expression]:
        """解析选择列表"""
        select_list = []

        if self.current_token and self.current_token.type == TokenType.STAR:
            select_list.append(ColumnReference("*"))
            self.advance()
        else:
            select_list.append(self.parse_expression())
            while self.current_token and self.current_token.type == TokenType.COMMA:
                self.advance()
                select_list.append(self.parse_expression())

        return select_list

    def parse_from_clause(self) -> FromClause:
        """解析FROM子句"""
        table_name = self.expect(TokenType.IDENTIFIER).value

        alias = None
        if self.current_token and self.current_token.type == TokenType.IDENTIFIER:
            alias = self.current_token.value
            self.advance()

        return FromClause(table_name, alias)

    def parse_insert_statement(self) -> InsertStatement:
        """解析INSERT语句"""
        self.expect(TokenType.INSERT)
        self.expect(TokenType.INTO)

        table_name = self.expect(TokenType.IDENTIFIER).value

        # 可选的列列表
        columns = None
        if self.current_token and self.current_token.type == TokenType.LPAREN:
            self.advance()
            columns = []
            columns.append(self.expect(TokenType.IDENTIFIER).value)
            while self.current_token and self.current_token.type == TokenType.COMMA:
                self.advance()
                columns.append(self.expect(TokenType.IDENTIFIER).value)
            self.expect(TokenType.RPAREN)

        self.expect(TokenType.VALUES)

        # 解析值列表
        values = []
        self.expect(TokenType.LPAREN)
        row_values = []
        row_values.append(self.parse_expression())
        while self.current_token and self.current_token.type == TokenType.COMMA:
            self.advance()
            row_values.append(self.parse_expression())
        self.expect(TokenType.RPAREN)
        values.append(row_values)

        return InsertStatement(table_name, columns, values)

    def parse_update_statement(self) -> UpdateStatement:
        """解析UPDATE语句"""
        self.expect(TokenType.UPDATE)

        table_name = self.expect(TokenType.IDENTIFIER).value

        self.expect(TokenType.SET)

        # 解析赋值列表
        assignments = []
        column_name = self.expect(TokenType.IDENTIFIER).value
        self.expect(TokenType.EQUALS)
        expression = self.parse_expression()
        assignments.append(Assignment(column_name, expression))

        while self.current_token and self.current_token.type == TokenType.COMMA:
            self.advance()
            column_name = self.expect(TokenType.IDENTIFIER).value
            self.expect(TokenType.EQUALS)
            expression = self.parse_expression()
            assignments.append(Assignment(column_name, expression))

        # 可选的WHERE子句
        where_clause = None
        if self.current_token and self.current_token.type == TokenType.WHERE:
            self.advance()
            where_clause = self.parse_expression()

        return UpdateStatement(table_name, assignments, where_clause)

    def parse_delete_statement(self) -> DeleteStatement:
        """解析DELETE语句"""
        self.expect(TokenType.DELETE)
        self.expect(TokenType.FROM)

        table_name = self.expect(TokenType.IDENTIFIER).value

        # 可选的WHERE子句
        where_clause = None
        if self.current_token and self.current_token.type == TokenType.WHERE:
            self.advance()
            where_clause = self.parse_expression()

        return DeleteStatement(table_name, where_clause)

    def parse_create_statement(self) -> CreateTableStatement:
        """解析CREATE TABLE语句"""
        self.expect(TokenType.CREATE)
        self.expect(TokenType.TABLE)

        table_name = self.expect(TokenType.IDENTIFIER).value

        self.expect(TokenType.LPAREN)

        # 解析列定义
        columns = []
        columns.append(self.parse_column_definition())

        while self.current_token and self.current_token.type == TokenType.COMMA:
            self.advance()
            columns.append(self.parse_column_definition())

        self.expect(TokenType.RPAREN)

        return CreateTableStatement(table_name, columns)

    def parse_drop_statement(self) -> DropTableStatement:
        """解析DROP TABLE语句"""
        self.expect(TokenType.DROP)
        self.expect(TokenType.TABLE)

        table_name = self.expect(TokenType.IDENTIFIER).value

        return DropTableStatement(table_name)

    def parse_column_definition(self) -> ColumnDefinition:
        """解析列定义"""
        column_name = self.expect(TokenType.IDENTIFIER).value

        # 解析数据类型
        data_type = self.parse_data_type()

        # 可选的约束
        constraints = []
        while self.current_token and (self.current_token.type == TokenType.IDENTIFIER or
                                     self.current_token.type in [TokenType.PRIMARY, TokenType.KEY,
                                                               TokenType.NOT, TokenType.NULL]):
            constraint = self.current_token.value.upper()
            if constraint in ['PRIMARY', 'NOT', 'UNIQUE', 'AUTO_INCREMENT'] or \
               self.current_token.type in [TokenType.PRIMARY, TokenType.NOT]:
                constraints.append(constraint)
                self.advance()
                # 处理复合约束如 PRIMARY KEY, NOT NULL
                if constraint == 'PRIMARY' and self.current_token and self.current_token.type == TokenType.KEY:
                    constraints[-1] = 'PRIMARY KEY'
                    self.advance()
                elif constraint == 'NOT' and self.current_token and self.current_token.type == TokenType.NULL:
                    constraints[-1] = 'NOT NULL'
                    self.advance()
            else:
                break

        return ColumnDefinition(column_name, data_type, constraints)

    def parse_data_type(self) -> str:
        """解析数据类型"""
        if self.current_token.type in [TokenType.INT, TokenType.VARCHAR, TokenType.FLOAT,
                                     TokenType.BOOLEAN, TokenType.DATE, TokenType.TEXT]:
            data_type = self.current_token.value
            self.advance()

            # 处理VARCHAR(n)的情况
            if data_type.upper() == 'VARCHAR' and self.current_token and self.current_token.type == TokenType.LPAREN:
                self.advance()
                size = self.expect(TokenType.NUMBER_LITERAL).value
                self.expect(TokenType.RPAREN)
                data_type = f"VARCHAR({size})"

            return data_type
        else:
            raise SyntaxError(f"Expected data type, but got {self.current_token.type.value}")

    def parse_expression(self) -> Expression:
        """解析表达式"""
        return self.parse_or_expression()

    def parse_or_expression(self) -> Expression:
        """解析OR表达式"""
        left = self.parse_and_expression()

        while self.current_token and self.current_token.type == TokenType.OR:
            operator = self.current_token.value
            self.advance()
            right = self.parse_and_expression()
            left = BinaryExpression(left, operator, right)

        return left

    def parse_and_expression(self) -> Expression:
        """解析AND表达式"""
        left = self.parse_not_expression()

        while self.current_token and self.current_token.type == TokenType.AND:
            operator = self.current_token.value
            self.advance()
            right = self.parse_not_expression()
            left = BinaryExpression(left, operator, right)

        return left

    def parse_not_expression(self) -> Expression:
        """解析NOT表达式"""
        if self.current_token and self.current_token.type == TokenType.NOT:
            self.advance()
            operand = self.parse_comparison_expression()
            return UnaryExpression("NOT", operand)
        else:
            return self.parse_comparison_expression()

    def parse_comparison_expression(self) -> Expression:
        """解析比较表达式"""
        left = self.parse_additive_expression()

        if self.current_token and self.current_token.type in [TokenType.EQUALS, TokenType.NOT_EQUALS,
                                                            TokenType.GREATER, TokenType.LESS,
                                                            TokenType.GREATER_EQUALS, TokenType.LESS_EQUALS,
                                                            TokenType.LIKE]:
            operator = self.current_token.value
            self.advance()
            right = self.parse_additive_expression()
            return BinaryExpression(left, operator, right)

        return left

    def parse_additive_expression(self) -> Expression:
        """解析加法表达式"""
        left = self.parse_multiplicative_expression()

        while self.current_token and self.current_token.value in ['+', '-']:
            operator = self.current_token.value
            self.advance()
            right = self.parse_multiplicative_expression()
            left = BinaryExpression(left, operator, right)

        return left

    def parse_multiplicative_expression(self) -> Expression:
        """解析乘法表达式"""
        left = self.parse_primary_expression()

        while self.current_token and self.current_token.value in ['*', '/']:
            operator = self.current_token.value
            self.advance()
            right = self.parse_primary_expression()
            left = BinaryExpression(left, operator, right)

        return left

    def parse_primary_expression(self) -> Expression:
        """解析基本表达式"""
        if not self.current_token:
            raise SyntaxError("Unexpected end of input")

        if self.current_token.type == TokenType.IDENTIFIER:
            identifier = self.current_token.value
            self.advance()

            # 检查是否是函数调用
            if self.current_token and self.current_token.type == TokenType.LPAREN:
                self.advance()
                arguments = []
                if self.current_token and self.current_token.type != TokenType.RPAREN:
                    arguments.append(self.parse_expression())
                    while self.current_token and self.current_token.type == TokenType.COMMA:
                        self.advance()
                        arguments.append(self.parse_expression())
                self.expect(TokenType.RPAREN)
                return FunctionCall(identifier, arguments)
            else:
                return ColumnReference(identifier)

        elif self.current_token.type == TokenType.STRING_LITERAL:
            value = self.current_token.value
            self.advance()
            return Literal(value, "STRING")

        elif self.current_token.type == TokenType.NUMBER_LITERAL:
            value = self.current_token.value
            self.advance()
            # 尝试转换为数字
            try:
                if '.' in value:
                    return Literal(float(value), "FLOAT")
                else:
                    return Literal(int(value), "INT")
            except ValueError:
                return Literal(value, "STRING")

        elif self.current_token.type == TokenType.LPAREN:
            self.advance()
            expression = self.parse_expression()
            self.expect(TokenType.RPAREN)
            return expression

        else:
            raise SyntaxError(f"Unexpected token in expression: {self.current_token.type.value}")


def print_ast(node: ASTNode, indent: int = 0):
    """打印抽象语法树"""
    prefix = "  " * indent

    if isinstance(node, SelectStatement):
        print(f"{prefix}SelectStatement:")
        print(f"{prefix}  Select List:")
        for expr in node.select_list:
            print_ast(expr, indent + 2)
        print_ast(node.from_clause, indent + 1)
        if node.where_clause:
            print(f"{prefix}  Where:")
            print_ast(node.where_clause, indent + 2)

    elif isinstance(node, FromClause):
        print(f"{prefix}From: {node.table_name}")
        if node.alias:
            print(f"{prefix}  Alias: {node.alias}")

    elif isinstance(node, BinaryExpression):
        print(f"{prefix}BinaryExpression ({node.operator}):")
        print_ast(node.left, indent + 1)
        print_ast(node.right, indent + 1)

    elif isinstance(node, ColumnReference):
        print(f"{prefix}Column: {node.column_name}")

    elif isinstance(node, Literal):
        print(f"{prefix}Literal: {node.value} ({node.value_type})")

    elif isinstance(node, CreateTableStatement):
        print(f"{prefix}CreateTable: {node.table_name}")
        for col in node.columns:
            print_ast(col, indent + 1)

    elif isinstance(node, ColumnDefinition):
        constraints_str = f" ({', '.join(node.constraints)})" if node.constraints else ""
        print(f"{prefix}Column: {node.column_name} {node.data_type}{constraints_str}")

    elif isinstance(node, InsertStatement):
        print(f"{prefix}InsertStatement:")
        print(f"{prefix}  Table: {node.table_name}")
        if node.columns:
            print(f"{prefix}  Columns: {', '.join(node.columns)}")
        print(f"{prefix}  Values:")
        for row in node.values:
            print(f"{prefix}    Row:")
            for expr in row:
                print_ast(expr, indent + 3)

    else:
        print(f"{prefix}{type(node).__name__}")


def test_parser():
    """测试语法分析器"""
    sql = """
    CREATE TABLE users (
        id INT PRIMARY KEY,
        name VARCHAR(50) NOT NULL,
        age INT,
        email VARCHAR(100)
    );

    INSERT INTO users (id, name, age, email) VALUES (1, 'Alice', 25, 'alice@example.com');

    SELECT id, name FROM users WHERE age > 20;
    """

    lexer = Lexer(sql)
    tokens = lexer.tokenize()

    parser = Parser(tokens)
    statements = parser.parse()

    print("语法分析结果 (AST):")
    for i, stmt in enumerate(statements):
        print(f"\n语句 {i + 1}:")
        print_ast(stmt)


if __name__ == "__main__":
    test_parser()
